#![allow(unused)]
use std::collections::{HashMap, HashSet,BTreeMap};
use std::fs::File;
use std::io::{self, Read};
use serde_json::Value;

// 定义并查集结构
struct UnionFind {
    parent: HashMap<String, String>, // 城市的父节点
}

impl UnionFind {
    /// 初始化并查集
    fn new() -> Self {
        UnionFind {
            parent: HashMap::new(),
        }
    }

    /// 查找根节点（路径压缩优化）
    fn find(&mut self, city: &String) -> String {
        if let Some(parent_city) = self.parent.get(city) {
            if parent_city != city {
                // 提前复制父节点，避免后续借用冲突
                let parent_city_cloned = parent_city.clone();
                let root = self.find(&parent_city_cloned);
                self.parent.insert(city.clone(), root.clone());
                return root;
            }
        }
        city.clone()
    }

    /// 合并两个城市
    fn union(&mut self, city1: &String, city2: &String) {
        let root1 = self.find(city1);
        let root2 = self.find(city2);

        if root1 != root2 {
            self.parent.insert(root1.clone(), root2);
        }
    }

    /// 添加城市到并查集
    fn add(&mut self, city: &String) {
        self.parent.entry(city.clone()).or_insert_with(|| city.clone());
    }
}

/// 计算每个批次的省份数
pub fn count_provinces() -> String {
    let district_data = match read_district_data() {
        Ok(data) => data,
        Err(_) => {
            println!("Error reading district data");
            return String::new();
        }
    };

    let mut result = String::new();

    // 遍历每个批次的数据
    for (batch_id, city_groups) in district_data.iter() {
        let mut uf: UnionFind = UnionFind::new();

        // 打印当前批次ID
        println!("Processing Batch ID: {}", batch_id);

        // 先去除重复城市键
        let mut cities_processed = HashSet::new();

        // 遍历每个城市和其邻居
        for (city, neighbors) in city_groups.iter() {
            if !cities_processed.contains(city) {
                uf.add(city);
                cities_processed.insert(city);  // 标记城市已经处理

                // 打印城市和邻居
                println!("City: {}, Neighbors: {:?}", city, neighbors);

                for neighbor in neighbors {
                    if !cities_processed.contains(neighbor) {
                        uf.add(neighbor);
                        cities_processed.insert(neighbor); // 标记邻居已处理
                    }
                    uf.union(city, neighbor); // 合并城市
                    println!("Union: {} -> {}", city, neighbor); // 打印合并操作
                }
            }
        }

        // 找到所有连通分量的根节点
        let mut province_roots = HashSet::new();
        for city in city_groups.keys() {
            let root = uf.find(city);
            println!("City: {}, Root: {}", city, root); // 打印根节点
            province_roots.insert(root);
        }

        // 批次结果
        println!("Batch ID: {}, Province Count: {}", batch_id, province_roots.len());
        result.push_str(&format!(
            "Batch ID: {}, Province Count: {}\n",
            batch_id,
            province_roots.len()
        ));
    }

    // 确保结果字符串末尾有换行符
    result.push_str("\n");

    // 获取批次数量
    let batch_count = result
        .lines()
        .filter(|line| line.starts_with("Batch ID:"))
        .count();

    // 打印批次数量
    println!("Total number of batches: {}", batch_count);

    // 处理结果为逗号分隔的格式
    let formatted_result = (1..=batch_count)
        .map(|batch_id| {
            result
                .lines()
                .find(|line| line.contains(&format!("Batch ID: {}", batch_id)))
                .and_then(|line| line.split("Province Count: ").nth(1))
                .unwrap_or("0")
                .trim()
                .to_string()
        })
        .collect::<Vec<_>>()
        .join(",");

    // 打印最终的逗号分隔格式结果
    println!("Formatted result: {}", formatted_result);

    formatted_result
}




/// 定义城市群数据结构
pub type DistrictData = BTreeMap<String, BTreeMap<String, Vec<String>>>;

pub fn read_district_data() -> Result<DistrictData, io::Error> {
    let mut file = File::open("topic1/solutiont2/district.json")?;
    let mut content = String::new();
    file.read_to_string(&mut content)?;

    // 使用 serde_json 自动反序列化 JSON
    let json_data: HashMap<String, HashMap<String, Vec<String>>> =
        serde_json::from_str(&content).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

    // 转换为 BTreeMap
    let district_data: DistrictData = json_data
    .into_iter()
    .map(|(batch, cities)| {
        let sorted_cities: BTreeMap<String, Vec<String>> = cities
            .into_iter()
            .map(|(city, neighbors)| {
                // 去重邻居列表
                let unique_neighbors = neighbors.into_iter().collect::<HashSet<_>>().into_iter().collect();
                (city, unique_neighbors)
            })
            .collect();
        (batch, sorted_cities)
    })
    .collect();

    Ok(district_data)
}